const fs = require('fs');
const path = require('path');

/**
 * Represents a Message entity.
 */
class Message {
    /**
     * Creates an instance of Message.
     * 
     * @param {string} role - The role of the message.
     * @param {string} content - The content of the message.
     * @param {string|null} [condition=null] - Optional condition for the message.
     */
    constructor(role, content, condition = null) {
        this.role = role;
        this.content = content;
        this.condition = condition;
    }

    /**
     * Converts the instance to an object.
     * 
     * @returns {Object} The message as an object.
     */
    toDict() {
        const dict = {
            role: this.role,
            content: this.content
        };

        if (this.condition !== null) {
            dict.condition = this.condition;
        }

        return dict;
    }

    /**
     * Returns a string representation of the message.
     * 
     * @returns {string} A string representation of the instance.
     */
    toString() {
        return `Message(role=${this.role}, content=${this.content}, condition=${this.condition})`;
    }
}


class MissingVariables extends Error {
    /**
     * Constructs a MissingVariables instance.
     * 
     * @param {string[]} variables - The list of missing variables.
     * @param {string|null} id - The identifier of the prompt.
     */
    constructor(variables, id) {
        super(`Missing variables ${variables} for prompt "${id}"`);
        this.variables = variables;
        this.id = id;
    }
}


class Value {
    /**
     * Constructs a Value instance.
     * 
     * @param {string} type - The type of the value.
     * @param {any} value - The value.
     */
    constructor(type, value) {
        this.type = type;
        this.value = value;
    }

    static Type = {
        ANY: "any",
        // Add other types as necessary
    };
}


/**
 * Represents a Prompt.
 */
class Prompt {
    /**
     * Creates an instance of Prompt.
     * 
     * @param {Array<Message>|Message|string} messages - The messages of the prompt.
     * @param {Object|null} [variables=null] - Optional variables for the prompt.
     * @param {Object|null} [config=null] - Optional configuration for the prompt.
     * @param {string|null} [id=null] - Optional ID for the prompt.
     * @param {HashMap<string, Value>|null} [variableDeclarations=null] - Optional variable declarations for the prompt.
     * @param {BaseClient|null} [client=null] - Optional client for the prompt.
     */
    constructor(messages, variables = null, config = null, id = null, variableDeclarations = null, client = null) {
        this.id = id;
        this.variables = variables;
        this.config = config;
        this.variableDeclarations = variableDeclarations;
        this._client = client;

        if (typeof messages === 'string') {
            this.messages = [new Message('user', messages)];
        } else if (messages instanceof Message) {
            this.messages = [messages];
        } else if (Array.isArray(messages)) {
            this.messages = messages.map(msg => {
                if (msg instanceof Message) {
                    return msg;
                } else if (typeof msg === 'object' && msg.role && msg.content) {
                    return new Message(msg.role, msg.content);
                } else {
                    throw new Error(`Invalid type for messages: ${typeof msg}`);
                }
            });
        } else {
            throw new Error(`Invalid type for messages: ${typeof messages}`);
        }
    }

    /**
     * Converts the instance to an object.
     * 
     * @returns {Object} The prompt as an object.
     */
    toDict() {
        return {
            messages: this.messages.map(msg => msg.toDict()),
            variables: this.variables,
            config: this.config,
            id: this.id,
        };
    }


    static _read(lines, id = null, client = null) {
        // Todo: add config parsing
        let config = {};
        let messages = [];
        let variableDeclarations = {};

        let currentMinSpaces = 0;
        let currentSection = null;
        let currentMessage = [];
        let condition = null;

        const addMessage = () => {
            if (currentSection) {
                if (currentSection != "variables" && !currentSection.startsWith("/") && currentSection != "config") {
                    messages.push(new Message(
                        currentSection,
                        currentMessage.map(m => m.slice(currentMinSpaces)).join("\n"),
                        condition
                    ));
                }
                currentMessage = [];
                currentMinSpaces = 0;
                condition = null;
            }
        };

        lines.split("\n").forEach(line => {
            // line = line.trimRight("\r");
            if (line.startsWith("<")) {
                line = line.trim();
                addMessage();
                currentSection = line.slice(1, -1);
                let ifIndex = currentSection.indexOf("if ");
                if (ifIndex !== -1) {
                    condition = currentSection.slice(ifIndex + 3);
                    currentSection = currentSection.slice(0, ifIndex).trim();
                } else {
                    condition = null;
                }
            } else if (currentSection === "config" && line.includes("=")) {
                let [key, value] = line.split("=", 1).map(s => s.trim());
                config[key] = value;
            } else if (currentSection === "variables") {
                let varValue = null;
                if (line.includes("=")) {
                    let split = line.split("=", 2);
                    line = split[0].trim();
                    varValue = split[1].trim();
                }
                
                let varType = Value.Type.ANY;
                if (line.includes(":")) {
                    let split = line.split(":", 2);
                    line = split[0].trim();
                    varType = split[1].trim();
                }

                let key = line.trim();
                variableDeclarations[key] = new Value(varType, varValue);
            } else if (["system", "user", "assistant", "if", "else"].includes(currentSection) || currentSection?.includes("add ")) {
                let minSpaces = line.length - line.trimLeft().length;
                if (minSpaces > 0 && (minSpaces < currentMinSpaces || currentMinSpaces === 0)) {
                    currentMinSpaces = minSpaces;
                }
                currentMessage.push(line);
            }
        });

        addMessage();
        return new this(messages, config, null, id, variableDeclarations, client);
    }

    static read(filePath, name = null) {
        let data = fs.readFileSync(filePath, 'utf8');
        if (name !== null) {
            let prompts = this.readAll(filePath);
            return prompts[name];
        } else {
            return this._read(data, name);
        }
    }

    static reads(text, name = null) {
        if (name !== null) {
            let prompts = this.readAll(texts);
            return prompts[name];
        } else {
            return this._read(text, name);
        }
    }

    static readAll(filePath, client = null) {
        let data = fs.readFileSync(filePath, 'utf8');
        return this.readsAll(data, client);
    }

    static readsAll(text, client = null) {
        let prompts = {};
        let lines = [];
        let currentPromptName = null;
        let currentMinSpaces = -1;

        text.split("\n").forEach(line => {
            line = line.trimRight();

            if (line.trimLeft().startsWith("<")) {
                let minSpaces = line.length - line.trimLeft().length;
                let strippedLine = line.trim();

                if (strippedLine.startsWith("<prompt") && minSpaces === 0) {
                    if (currentPromptName) {
                        prompts[currentPromptName] = this._read(lines.map(m => m.slice(currentMinSpaces)).join("\n"), currentPromptName, client);
                    }
                    currentMinSpaces = -1;
                    lines = [];
                    let tempPromptName = strippedLine.slice(8, -1).trim();
                    let id = null;
                    if (tempPromptName.includes("id=")) {
                        [tempPromptName, id] = tempPromptName.split("id=").map(s => s.trim().replace(/^["']|["']$/g, ''));
                        currentPromptName = id;
                    }
                } else if (strippedLine.startsWith("</prompt>") && minSpaces === 0) {
                    prompts[currentPromptName] = this._read(lines.map(m => m.slice(currentMinSpaces)).join("\n"), currentPromptName, client);
                    currentPromptName = null;
                    currentMinSpaces = -1;
                    lines = [];
                } else {
                    lines.push(line);
                    if (currentMinSpaces === -1 || minSpaces < currentMinSpaces) {
                        currentMinSpaces = minSpaces;
                    }
                }
            } else {
                lines.push(line);
            }
        });

        return prompts;
    }

    _file() {
        return this.messages.map(message => `<${message.role.toLowerCase()}>\n    ${message.content.split("\n").join("\n    ")}`).join("\n");
    }

    static write(prompt, filePath) {
        let content = '';
        if (typeof prompt === 'object' && !Array.isArray(prompt)) {
            for (let [name, p] of Object.entries(prompt)) {
                content += `<prompt id="${name}">\n    ${p._file().split("\n").join("\n    ")}\n</prompt>\n\n`;
            }
        } else {
            content = prompt._file();
        }

        fs.writeFileSync(filePath, content);
    }

    static fromOpenAI(messages) {
        return new this(
            messages.map(message => new Message(message.role, message.content))
        );
    }

    _checkCondition(condition) {
        if (!condition) {
            return false;
        }

        condition = condition.trim();
        if (condition.startsWith('not')) {
            return !this._checkCondition(condition.slice(3).trim());
        }
        if (condition.startsWith('!')) {
            return !this._checkCondition(condition.slice(1).trim());
        }

        let varName = condition.trim().replace(/[{}]/g, '');
        if (!(varName in this.variables)) {
            return false;
        }

        let value = this.variables[varName];
        return value === 'true' || value === 'True' || value === '1' || value === 1 || value === true;
    }

    static _formatMap(text, variables) {
        let pattern = /\{[^{}]*\}/g;
        let matches = Array.from(text.matchAll(pattern));
        for (let i = matches.length - 1; i >= 0; i--) {
            let match = matches[i];
            let varText = match[0].slice(1, -1);
            if (varText in variables) {
                text = text.slice(0, match.index) + variables[varText] + text.slice(match.index + match[0].length);
            }
        }
        return text;
    }

    _convertExpression(expression) {
        expression = expression.trim();

        if (expression.startsWith('"') && expression.endsWith('"')) {
            return Prompt._formatMap(
                expression.slice(1, -1).replace(/\\n/g, "\n").replace(/\\t/g, "\t"),
                this.variables
            );
        } else {
            return Prompt._formatMap('{' + expression + '}', this.variables);
        }
    }

    _formatMessage(message, overrideRole = null) {
        let text = message.content;
        let pattern = /\{[^{}]*\}/g;
        let matches = Array.from(text.matchAll(pattern));

        for (let i = matches.length - 1; i >= 0; i--) {
            let match = matches[i];
            let varText = match[0].slice(1, -1);

            if (varText.includes("?")) {
                let [condition, values] = varText.split("?");
                let [left, right] = values.split(":");
                if (this._checkCondition(condition)) {
                    text = text.slice(0, match.index) + this._convertExpression(left) + text.slice(match.index + match[0].length);
                } else {
                    text = text.slice(0, match.index) + this._convertExpression(right) + text.slice(match.index + match[0].length);
                }
            }
        }

        return {
            role: overrideRole !== null ? overrideRole : message.role,
            content: Prompt._formatMap(text, this.variables),
        };
    }

    compose(overridePrompt = null) {
        let requiredVariables = new Set(
            Object.entries(this.variableDeclarations)
                .filter(([_, value]) => value.value === null)
                .map(([varName, _]) => varName)
        );
        let definedVariables = new Set(Object.keys(this.variables || {}));
    
        let prevRole = null, prevCondition = false;
        let ret = [];
    
        const addPrompt = name => {
            if (this._client && name in this._client.prompts) {
                for (let subMsg of this.compose(this._client.prompts[name])) {
                    if (subMsg.role === "system") {
                        throw new Error("System prompt not supported yet.");
                    } else {
                        ret.push(subMsg);
                    }
                }
            } else {
                throw new Error(`Prompt ${name} not found.`);
            }
        };
    
        let messages = overridePrompt === null ? this.messages : overridePrompt.messages;
        for (let message of messages) {
            let add = false;
    
            if (message.condition !== null) {
                add = this._checkCondition(message.condition);
                prevRole = "if";
                prevCondition = add;
            } else if (message.role === "else") {
                if (!prevCondition) {
                    if (message.role.startsWith("else add ")) {
                        let id = message.role.substring(9).trim();
                        if (id.includes("id=")) {
                            id = id.split("id=")[1].trim().replace(/^["']|["']$/g, '');
                        }
                        addPrompt(id);
                    } else {
                        ret.push(this._formatMessage(message, prevRole));
                    }
                }
                prevRole = "else";
                prevCondition = true;
                continue;
            } else {
                add = true;
                prevRole = message.role;
            }
    
            if (add) {
                if (message.role.startsWith("add ")) {
                    let id = message.role.substring(4).trim();
                    if (id.includes("id=")) {
                        id = id.split("id=")[1].trim().replace(/^["']|["']$/g, '');
                    }
                    addPrompt(id);
                } else {
                    ret.push(this._formatMessage(message));
                }
            }
        }
    
        if ([...requiredVariables].filter(varName => !definedVariables.has(varName)).length > 0) {
            throw new MissingVariables(
                [...requiredVariables].filter(varName => !definedVariables.has(varName)), 
                overridePrompt ? overridePrompt.id : this.id
            );
        }
    
        return ret;
    }
    

    static _applyVariables(messages, variables) {
        return messages.map(message => new Message(
            message.role,
            Prompt._formatMap(message.content, variables),
            message.condition
        ));
    }

    _checkDuplicateKeys(otherVariables) {
        let duplicateKeys = new Set(Object.keys(this.variables || {}));
        for (let key of Object.keys(otherVariables || {})) {
            if (duplicateKeys.has(key) && this.variables[key] !== otherVariables[key]) {
                duplicateKeys.delete(key);
            }
        }

        let result = {};
        for (let key of duplicateKeys) {
            result[key] = this.variables[key];
        }

        return result;
    }

    _removeDuplicateKeysFromMessages(otherVariables) {
        let messages = this.messages;
        let appliedVariables = this._checkDuplicateKeys(otherVariables);
        if (Object.keys(appliedVariables).length > 0) {
            messages = Prompt._applyVariables(this.messages, appliedVariables);
        }

        return messages;
    }

    format(...args) {
        let kwargs = Object.assign({}, ...args);
        let messages = this._removeDuplicateKeysFromMessages(kwargs);

        return new Prompt(
            messages.map(message => new Message(message.role, message.content, message.condition)),
            {...(this.variables || {}), ...kwargs},
            this.config,
            this.id,
            this.variableDeclarations,
            this._client
        );
    }

    add(other) {
        if (other instanceof Message) {
            return new Prompt(
                [...this.messages, other],
                {...(this.variables || {})}
            );
        } else if (other instanceof Prompt) {
            let messages = this._removeDuplicateKeysFromMessages(other.variables || {});
            return new Prompt(
                [...messages, ...other.messages],
                {...(this.variables || {}), ...(other.variables || {})}
            );
        } else {
            throw new Error('NotImplementedError');
        }
    }

    toString() {
        let messageStr = this.messages.map(message => `${message.role}: ${message.content}`).join('\n');
        return `Prompt(\n${messageStr}\n${JSON.stringify(this.variables || {}, null, 2)}\n)`;
    }
}

class Response {
    /**
     * Constructs a Response instance.
     * 
     * @param {string} content - The content of the response.
     * @param {number|null} [promptTokens=null] - The number of tokens in the prompt.
     * @param {number|null} [completionTokens=null] - The number of tokens in the completion.
     * @param {Object|null} [raw=null] - The raw response data.
     * @param {Object} [additionalProperties={}] - Additional properties for the response.
     */
    constructor(content, promptTokens = null, completionTokens = null, raw = null, additionalProperties = {}) {
        this.content = content;
        this.promptTokens = promptTokens;
        this.completionTokens = completionTokens;
        this.raw = raw;

        // Handling additional properties similar to **kwargs in Python
        for (const [key, value] of Object.entries(additionalProperties)) {
            this[key] = value;
        }
    }

    /**
     * Creates a new Response instance from various response types.
     * 
     * @param {Response|string} response - The response to create a Response instance from.
     * @returns {Response} A new Response instance.
     */
    static create(response) {
        if (response instanceof Response) {
            return response;
        } else if (typeof response === 'string') {
            return new Response(response);
        } else {
            throw new Error('NotImplementedError');
        }
    }

    /**
     * Returns a string representation of the Response instance.
     * 
     * @returns {string} String representation of the instance.
     */
    toString() {
        return `Response(content=${this.content}, promptTokens=${this.promptTokens}, completionTokens=${this.completionTokens}, raw=${JSON.stringify(this.raw)})`;
    }

    /**
     * Converts the instance to an object.
     * 
     * @returns {Object} The response as an object.
     */
    toDict() {
        return {
            content: this.content,
            prompt_tokens: this.promptTokens,
            completion_tokens: this.completionTokens,
            raw: this.raw,
        };
    }
}

module.exports = {
    Message,
    Prompt,
    Response
};
